from lida import Manager,TextGenerationConfig,llm



#from llmx import llm
model_name = "moonshot-v1-8k"


text_gen = llm(provider="openai", api_base="https://api.moonshot.cn/v1", api_key=API_KEY, model=model_name)

textgen_config = TextGenerationConfig(n=1, temperature=0.5, model=model_name, use_cache=True)


lida = Manager(text_gen)

summary = lida.summarize("../data/cars.csv", summary_method="default", textgen_config=textgen_config)  
print(summary)


goals = lida.goals(summary, n=2, textgen_config=textgen_config)

for goal in goals:
    print(goal)

library = "seaborn"
user_query = "What is the average price of cars by type?"
instructions = ["make the chart height and width equal", "change the color of the chart to red"]

charts = lida.visualize(
    summary,
    goal=user_query,
    library=library,
    textgen_config=textgen_config,
)
charts[0].savefig('example_18.png')

code = charts[0].code

edited_charts = lida.edit(code=code, summary=summary, instructions=instructions, library=library, textgen_config=textgen_config)
edited_charts[0].savefig('example_18_edited.png')

explanation = lida.explain(code, textgen_config=textgen_config)
for row in explanation[0]:
    print(row["section"], '   ', row["explanation"])

evaluations = lida.evaluate(code=code, goal=goals[0], library=library, textgen_config=textgen_config)
print(evaluations)

recommended_charts =  lida.recommend(code=code, summary=summary, n=2,  textgen_config=textgen_config)
recommended_charts[0].savefig('example_18_recommended.png')